import ast
import os
from typing import List, Tuple, Optional


def extract_imports(file_path: str) -> List[Tuple]:
    """
    Parse a Python file and extract all imported modules.

    Args:
        file_path: Path to the Python file to parse

    Returns:
        A list of tuples containing import information:
        - For regular imports: (module_name, 0)
        - For from imports: (module_name, level)
        - For from . import y: (None, level)
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File {file_path} not found")

    with open(file_path, "r") as file:
        source = file.read()

    # Parse the source code into an AST
    tree = ast.parse(source)

    imported_modules = []

    # Walk through all nodes in the AST
    for node in ast.walk(tree):
        # Handle standard imports (import x, import x.y)
        if isinstance(node, ast.Import):
            for name in node.names:
                imported_modules.append((name.name, 0))

        # Handle from imports (from x import y)
        elif isinstance(node, ast.ImportFrom):
            if node.module:  # Handles 'from x import y'
                imported_modules.append((node.module, node.level))
            else:  # Handles 'from . import y'
                imported_modules.append((None, node.level))

    return imported_modules


def find_module_path(
    module_name: str,
    python_paths: List[str],
    base_dir: Optional[str] = None,
    level: int = 0,
) -> Optional[str]:
    """
    Find the file path for a given module name.

    Args:
        module_name: Name of the module to find (e.g., 'a.b.c')
        python_paths: List of paths to search in
        base_dir: Base directory for relative imports
        level: Level of relative import (e.g., from ..a import b)

    Returns:
        Path to the module file if found, None otherwise
    """
    if level > 0 and base_dir:
        # Handle relative imports
        rel_path = base_dir
        for _ in range(level - 1):
            rel_path = os.path.dirname(rel_path)

        # Make sure rel_path is within one of the python_paths
        in_python_path = False
        for python_path in python_paths:
            abs_python_path = os.path.abspath(python_path)
            abs_rel_path = os.path.abspath(rel_path)
            if abs_rel_path.startswith(abs_python_path):
                in_python_path = True
                break

        if not in_python_path:
            return None

        if module_name:
            path_to_check = os.path.join(rel_path, *module_name.split("."))
        else:
            path_to_check = rel_path
    else:
        # Handle absolute imports
        for python_path in python_paths:
            path_parts = module_name.split(".")
            path_to_check = os.path.join(python_path, *path_parts)

            # Check if it's a file
            if os.path.isfile(f"{path_to_check}.py"):
                return f"{path_to_check}.py"

            # Check if it's a directory with __init__.py
            if os.path.isdir(path_to_check) and os.path.isfile(
                os.path.join(path_to_check, "__init__.py")
            ):
                return os.path.join(path_to_check, "__init__.py")

        return None  # Not found in any of the python_paths

    # Check if it's a file
    if os.path.isfile(f"{path_to_check}.py"):
        return f"{path_to_check}.py"

    # Check if it's a directory with __init__.py
    if os.path.isdir(path_to_check) and os.path.isfile(
        os.path.join(path_to_check, "__init__.py")
    ):
        return os.path.join(path_to_check, "__init__.py")

    return None


def find_all_dependencies(file_path: str, python_paths: List[str]) -> List[str]:
    """
    Find all Python file dependencies using DFS.

    Args:
        file_path: Path to the Python file to start from
        python_paths: List of paths to search for imports

    Returns:
        A list of file paths that the input file depends on
    """
    visited = set()
    dependencies = []

    def dfs(current_file: str):
        if current_file in visited:
            return

        visited.add(current_file)
        current_dir = os.path.dirname(current_file)

        try:
            imports = extract_imports(current_file)

            for module_name, level in imports:
                if module_name:
                    # Handle regular imports (import x) and from imports (from x import y)
                    module_path = find_module_path(
                        module_name, python_paths, current_dir, level
                    )
                    if (
                        module_path
                        and module_path not in visited
                        and os.path.isfile(module_path)
                    ):
                        dependencies.append(module_path)
                        dfs(module_path)
                elif level > 0:
                    # Handle 'from . import y' cases
                    # Here we need to find all Python files in the directory level steps up
                    rel_path = current_dir
                    for _ in range(level):
                        rel_path = os.path.dirname(rel_path)

                    # Make sure rel_path is within one of the python_paths
                    in_python_path = False
                    for python_path in python_paths:
                        abs_python_path = os.path.abspath(python_path)
                        abs_rel_path = os.path.abspath(rel_path)
                        if abs_rel_path.startswith(abs_python_path):
                            in_python_path = True
                            break

                    if not in_python_path:
                        continue

                    # Check for __init__.py
                    init_path = os.path.join(rel_path, "__init__.py")
                    if os.path.isfile(init_path) and init_path not in visited:
                        dependencies.append(init_path)
                        dfs(init_path)

                    # Find all .py files in the directory
                    if os.path.isdir(rel_path):
                        for item in os.listdir(rel_path):
                            if item.endswith(".py") and item != "__init__.py":
                                py_path = os.path.join(rel_path, item)
                                if py_path not in visited and os.path.isfile(py_path):
                                    dependencies.append(py_path)
                                    dfs(py_path)

        except Exception as e:
            print(f"Error processing {current_file}: {e}")

    dfs(file_path)
    return dependencies


if __name__ == "__main__":
    # Example usage
    test_file = "ast_testpy.py"  # Replace with an actual file path
    python_paths = ["."]  # Add your PYTHONPATH directories here
    try:
        dependencies = find_all_dependencies(test_file, python_paths)
        print(f"Dependencies for {test_file}:")
        for dep in dependencies:
            print(f"  - {dep}")
    except FileNotFoundError as e:
        print(e)
